#!/bin/bash
#
# *** Managed by salt ***
#
# -*- coding: utf-8 -*-
# Author: Pär Andersson (National Supercomputer Centre, Sweden)
# Version: 0.3 2007-07-30
# Modified for use by slurm-formula
#
# This will submit a batch script that starts screen on a node.
# Then ssh is used to connect to the node and attach the screen.
# The result is very similar to an interactive shell in PBS
# (qsub -I)

# Batch Script that starts SCREEN
BS={{slurm.bindir}}/srun-x-start
# Interactive screen script
IS={{slurm.bindir}}/srun-x-session

# Submit the job and get the job id
JOB=`sbatch --output=/dev/null --error=/dev/null $@ $BS 2>&1 \
    | egrep -o -e "\b[0-9]+$"`

# Make sure the job is always canceled
trap "{ {{slurm.bindir}}/scancel -q $JOB; exit; }" SIGINT SIGTERM EXIT

echo "Waiting for JOBID $JOB to start"
while true;do
    sleep 1s

    # Check job status
    STATUS=`squeue -j $JOB -t PD,R -h -o %t`

    if [ "$STATUS" = "R" ];then
	# Job is running, break the while loop
	break
    elif [ "$STATUS" != "PD" ];then
	echo "Job is not Running or Pending. Aborting"
	scancel $JOB
	exit 1
    fi

    echo -n "."

done

# Determine the first node in the job:
NODE=`srun --jobid=$JOB -N1 hostname`

# SSH to the node and attach the screen
sleep 1s
ssh -X -t $NODE $IS slurm$JOB
# The trap will now cancel the job before exiting.
